#!/usr/bin/env python

import optparse
import os
import sys

import minuit
import numpy
import matplotlib
from matplotlib import pyplot
import tables

import histogram
from charmicat import BadCommandLineError

import ypipi

def gauss(x, A, mu, sigma):
    return A*numpy.exp(-0.5*((x - mu)/sigma)**2)

def bg(x, p0, p1, p2, p3, p4, p5, p6):
    return numpy.polynomial.chebyshev.chebval(x, [p0, p1, p2, p3, p4, p5, p6])

def y_calc(x, A, mu, sigma, p0, p1, p2, p3, p4, p5, p6):
    return gauss(x, A, mu, sigma) + bg(x, p0, p1, p2, p3, p4, p5, p6)

def chisq(A, mu, sigma, p0, p1, p2, p3, p4, p5, p6):
    global _g_x
    global _g_y
    global _g_yunc
    global _g_ncalls

    chi2 = (y_calc(_g_x, A, mu, sigma, p0, p1, p2, p3, p4, p5, p6) - _g_y) / _g_yunc

    if _g_ncalls % 100 == 0:
        print _g_ncalls, numpy.sum(chi2)

    _g_ncalls += 1

    return numpy.sum(chi2**2)

def bg_chisq(p0, p1, p2, p3, p4, p5, p6):
    global _g_x
    global _g_y
    global _g_yunc
    global _g_ncalls

    chi2 = (bg(_g_x, p0, p1, p2, p3, p4, p5, p6) - _g_y) / _g_yunc

    if _g_ncalls % 100 == 0:
        print 'bg_chisq:',_g_ncalls, numpy.sum(chi2)

    _g_ncalls += 1

    return numpy.sum(chi2**2)

def sig_chisq(A, mu, sigma):
    global _g_x
    global _g_y
    global _g_yunc
    global _g_ncalls

    chi2 = (gauss(_g_x, A, mu, sigma) - _g_y) / _g_yunc

    if _g_ncalls % 100 == 0:
        print 'sig_chisq:',_g_ncalls, numpy.sum(chi2)

    _g_ncalls += 1

    return numpy.sum(chi2**2)



def main():
    '''
    ypipi_fit_Y2S.py [opts] hists.hpkl

    Fit histogram `mm_off_pi0pi0' from 9.35 to 9.7 GeV with Gauss + PolN
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    opts, args = op.parse_args()

    if len(args) == 0:
        msg = 'Please supply an input file.'
        raise BadCommandLineError(msg, -1)

    hdict = histogram.load(args[0])
    
    h = hdict['mm_off_pi0pi0']

    global _g_x, _g_y, _g_yunc, _g_ncalls

    # Begin BG fit

    lower_limit        = 0
    signal_lower_limit = h.bin_index(9.4)
    signal_upper_limit = h.bin_index(9.6)
    upper_limit        = h.bin_index(9.7)

    y_bg = numpy.concatenate((h.bins()[lower_limit:signal_lower_limit], h.bins()[signal_upper_limit:upper_limit]))
    x_full = numpy.array(map(h.bin_center, range(h.nbins)))
    x_bg = numpy.concatenate((x_full[lower_limit:signal_lower_limit], x_full[signal_upper_limit:upper_limit]))

    yunc_bg = numpy.sqrt(y_bg)
    yunc_bg = numpy.where(yunc_bg > 0, yunc_bg, numpy.ones(len(yunc_bg)))

    _g_x    = x_bg
    _g_y    = y_bg
    _g_yunc = yunc_bg

    m_bg = minuit.Minuit(bg_chisq)

    m_bg.values['p0'] = 1
    m_bg.values['p1'] = 1
    m_bg.values['p2'] = 1
    m_bg.values['p3'] = 0
    m_bg.values['p4'] = 0
    m_bg.values['p5'] = 0
    m_bg.values['p6'] = 0

    m_bg.fixed['p4'] = True
    m_bg.fixed['p5'] = True
    m_bg.fixed['p6'] = True

    initial_values = {}
    initial_values.update(m_bg.values)

    _g_ncalls = 0
    
    try:
        m_bg.migrad()
    except minuit.MinuitError:
        raise
    
    
    figx = 10
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.6, 0.8))

    info = [m_bg.values['p0'], m_bg.errors['p0'],
            m_bg.values['p1'], m_bg.errors['p1'],
            m_bg.values['p2'], m_bg.errors['p2'],
            m_bg.values['p3'], m_bg.errors['p3'],
            m_bg.values['p4'], m_bg.errors['p4'],
            m_bg.values['p5'], m_bg.errors['p5'],
            m_bg.values['p6'], m_bg.errors['p6'],
            bg_chisq(**initial_values) / (len(_g_x) - len(m_bg.values.keys())), 
            bg_chisq(**m_bg.values) / (len(_g_x) - len(m_bg.values.keys())), 
            ]

    fptex = ypipi.textable.write_bg_table(info)

    pyplot.text(1.35, 0.75,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = bg(xfit, **m_bg.values)

    pyplot.plot(xfit, yfit, 'b-')
#    pyplot.plot(xfit, y_calc(xfit, **initial_values), 'k--')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title + r' [all4g $>$ 50 MeV]')

    path_out = h.name + '_bgfit.png'
    pyplot.savefig(path_out)
    path_out = h.name + '_bgfit.pdf'
    pyplot.savefig(path_out)


    _g_y = h.bins()[lower_limit:upper_limit]
    _g_x = x_full[lower_limit:upper_limit]
    _g_yunc = numpy.sqrt(_g_y)
    _g_yunc = numpy.where(_g_yunc > 0, _g_yunc, numpy.ones(len(_g_yunc)))

    m = minuit.Minuit(chisq)

    m.values['A']     = 1.89
    m.values['mu']    = 9.5
    m.values['sigma'] = 0.027

    for k in 'p0 p1 p2 p3 p4 p5 p6'.split():
        m.values[k] = m_bg.values[k]

        m.fixed[k] = True

    initial_values = {}
    initial_values.update(m.values)

    try:
        m.migrad()
    except minuit.MinuitError:
        raise

    figx = 10
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.6, 0.8))

    info = [m.values['A'], m.errors['A'],
            m.values['mu'], m.errors['mu'],
            m.values['sigma'], m.errors['sigma'],
            m.values['p0'], m.errors['p0'],
            m.values['p1'], m.errors['p1'],
            m.values['p2'], m.errors['p2'],
            m.values['p3'], m.errors['p3'],
            chisq(**initial_values) / (len(_g_x) - len(m.values.keys())), 
            chisq(**m.values) / (len(_g_x) - len(m.values.keys())), 
            ]

    fptex = ypipi.textable.write_table(info)

    pyplot.text(1.35, 0.75,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = y_calc(xfit, **m.values)

    pyplot.plot(xfit, yfit, 'b-')
#    pyplot.plot(xfit, y_calc(xfit, **initial_values), 'k--')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title + r' [all4g $>$ 50 MeV, bg fixed]')

    path_out = h.name + '_bgfixed_fit.png'
    pyplot.savefig(path_out)
    path_out = h.name + '_bgfixed_fit.pdf'
    pyplot.savefig(path_out)
    

    _g_y -= bg(_g_x, **m_bg.values)
    _g_yunc = numpy.sqrt(numpy.fabs(_g_y))
    _g_yunc = numpy.where(_g_yunc > 0, _g_yunc, numpy.ones(len(_g_yunc)))    


    m_sig = minuit.Minuit(sig_chisq)

    m_sig.values['A'] = m.values['A']
    m_sig.values['mu'] = m.values['mu']
    m_sig.values['sigma'] = m.values['sigma']

    initial_values = {}
    initial_values.update(m_sig.values)

    try:
        m_sig.migrad()
    except minuit.MinuitError:
        print "Died on subtracted fit"
        raise

    

    figx = 8
    figy = 6
    fig = pyplot.figure(figsize=(figx,figy))

    axes = pyplot.axes((0.135, 0.125, 0.8, 0.8))

    info = [m_sig.values['A'], m_sig.errors['A'],
            m_sig.values['mu'], m_sig.errors['mu'],
            m_sig.values['sigma'], m_sig.errors['sigma'],
            sig_chisq(**initial_values) / (len(_g_x) - len(m_sig.values.keys())), 
            sig_chisq(**m_sig.values) / (len(_g_x) - len(m_sig.values.keys())), 
            ]

    fptex = ypipi.textable.write_sig_table(info)

    pyplot.text(0.95, 0.95,
                "\input{" + fptex.name + "}",
                fontsize=18, verticalalignment='top',
                horizontalalignment='right',
                transform=axes.transAxes,
                bbox={'boxstyle':'round', 'ec':'k', 'color':'#FFFFFF', 'alpha': '0.7'})

    pyplot.errorbar(_g_x, _g_y, yerr=_g_yunc, fmt='r.')

    xfit = numpy.linspace(h.bin_center(lower_limit), 
                          h.bin_center(upper_limit), 
                          10000)

    yfit = gauss(xfit, **m_sig.values)

    pyplot.plot(xfit, yfit, 'b-')
#    pyplot.plot(xfit, y_calc(xfit, **initial_values), 'k--')

    pyplot.xlabel(r'MM($\pi^0\pi^0$) GeV')
    pyplot.ylabel(r'Counts per 5 MeV')

    pyplot.title(h.title + r' [all4g $>$ 50 MeV, bg subtracted]')

    path_out = h.name + '_subtracted_fit.png'
    pyplot.savefig(path_out)

    path_out = h.name + '_subtracted_fit.pdf'
    pyplot.savefig(path_out)
    


    return 0
    
    

if __name__ == "__main__":
    try:
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
